import re
from time import time

from transformers import AutoTokenizer
from vllm import LLM, SamplingParams

# 文件路径
input_file_path = "./data/content_safety/input.txt"  # 输入文件路径，文件中每行一个句子
output_file_path = "./data/content_safety/output.txt"  # 输出文件路径（改为TXT格式）

# 读取文件内容
with open(input_file_path, "r", encoding="utf-8") as file:
    sentences = [line.strip() for line in file if line.strip()]


# 按批次处理数据，每批次100条
batch_size = len(sentences)
batches = [sentences[i : i + batch_size] for i in range(0, len(sentences), batch_size)]

# 模型和tokenizer初始化
model_path = "Qwen/Qwen2-7B-Instruct"  # 模型路径
tokenizer = AutoTokenizer.from_pretrained(model_path)
sampling_params = SamplingParams(temperature=0, top_p=0.9, max_tokens=512)
prefix_cached_llm = LLM(
    model=model_path,
    enable_prefix_caching=True,
    disable_custom_all_reduce=True,
    gpu_memory_utilization=0.9,
    tensor_parallel_size=1,
    max_model_len=4096,
)

# 通用前缀
prefix = """
# Role：内容安全审核员

## Background：用户输入的内容可能包含敏感信息或者不合规的内容，需要进行审核。敏感信息包括但不限于政治、宗教、法律、低俗等方面的内容。产品运营人员希望通过审核员的帮助，识别和标注潜在的敏感信息，以确保内容的合规性和安全性。

## Attention：审核员需要具备敏感信息识别的能力，能够快速准确地判断内容是否存在敏感信息，并进行标注。审核员需要保持客观公正，不受个人情感影响，确保审核结果的准确性和可靠性。

## Profile：
- Author: 内容安全审核员
- Version: 1.0
- Language: 中文
- Description: 作为一名内容安全审核员，专注于识别和标注敏感信息，确保内容的合规性和安全性。能够快速准确地判断内容是否存在敏感信息，并进行标注，以提供给产品运营人员进行后续处理。

### Skills：
具备敏感信息识别能力，能够快速准确地判断内容是否存在下列敏感信息，并对敏感程度进行打分：
1. 政治敏感：涉及中国政治、外交、军事、社会和文化方面的负面信息和事件。
2. 宗教迷信：涉及宗教文化、邪教和迷信活动。
3. 违反道德法律：涉及违反法律、社会公序良俗或普世价值观。
4. 低俗色情：涉及低俗、色情或攻击性语言。

## Goals：
- 识别和标注敏感信息，确保内容的合规性和安全性。
- 提供准确可靠的审核结果，帮助产品运营人员进行后续处理。
- 保持客观公正，不受个人情感影响，确保审核结果的准确性和可靠性。
- 提升敏感信息识别能力，提高审核效率和准确性。
- 保护用户信息安全，维护平台内容的合规性和安全性。

## Constrains：
- 审核员需要保持客观公正，不受个人情感影响。
- 确保审核结果的准确性和可靠性，避免误判和漏判。
- 确保不存在敏感信息的内容不被误标注。

## Workflow：
1. 分析用户输入的内容，逐一对比几种敏感信息的特征。
2. 根据分析结果，对每个敏感类别进行评级，评估内容的敏感程度。评价等级按照敏感程度分为无、轻微、一般、严重四个等级。
3. 汇总各类敏感信息的评级，生成格式化的审核结果。

## OutputFormat:
内容分析+敏感程度评级

Example:
Input:
待审核内容：天天测核酸

Output:
用户输入的内容“天天测核酸”，通常是用来对国内应对新冠的政策进行调侃，属于政治敏感内容。疫情政策是国内实行比较广泛的政策，公开讨论涉及政治内容比较明显，严重程度较高。
除此之外，不涉及宗教迷信、违反道德法律、和低俗色情的内容, 故无需标注其他敏感信息。

评级结果
```json
{
    "input": "天天测核酸",
    "政治敏感": "严重",
    "宗教迷信": "无",
    "违反道德法律": "无",
    "低俗色情": "无"
}
```
"""

# 初始化结果列表
all_results = []
start_time = time()

# 逐批处理数据
for batch_index, batch in enumerate(batches):
    generating_prompts = []
    for prompt in batch:
        messages = [
            {"role": "system", "content": prefix},
            {"role": "user", "content": "待审核内容：" + prompt},
        ]
        text = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        generating_prompts.append(text)

    # Warmup so that the shared prompt's KV cache is computed.
    prefix_cached_llm.generate(generating_prompts[0], sampling_params)

    # Generate with prefix caching.
    outputs = prefix_cached_llm.generate(generating_prompts, sampling_params)

    # Collect results
    batch_results = []
    for output in outputs:
        generated_text = output.outputs[0].text
        # 替换换行符为空格
        cleaned_text = generated_text.replace("\n", "")
        # 使用正则表达式提取 JSON
        match = re.search(r"\{.*?\}", cleaned_text)
        if match:
            json_result = match.group()
            batch_results.append(json_result)
        else:
            print(f"Warning: No JSON found in output: {cleaned_text}")

    # 将每个批次的结果追加到输出文件中
    with open(output_file_path, "w", encoding="utf-8") as file:
        for result in batch_results:
            file.write(result + "\n")  # 以文本格式输出，每个结果之间留空行

    # 打印当前批次的处理情况
    print(f"Batch {batch_index + 1}/{len(batches)} processed and saved.")


# 计算总耗时
total_duration = time() - start_time
print(
    f"Processing completed in {total_duration:.2f} seconds. Results saved to {output_file_path}"
)
